import matplotlib.pyplot as plt
import pandas as pd
import yfinance as yf
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split

# Fetch data
data = yf.download('AAPL', start='2020-01-01', end='2021-01-01')
data['Date'] = data.index

# Feature engineering
data['Year'] = data['Date'].dt.year
data['Month'] = data['Date'].dt.month
data['Day'] = data['Date'].dt.day

# Prepare data
X = data[['Year', 'Month', 'Day']]
y = data['Close']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

# Model training
model = LinearRegression()
model.fit(X_train, y_train)

# Prediction
data['Prediction'] = model.predict(data[['Year', 'Month', 'Day']])

# Visualization
plt.figure(figsize=(5, 4))
plt.plot(data['Date'], data['Close'], label='True')
plt.plot(data['Date'], data['Prediction'], label='Prediction')
plt.legend()
plt.show()

